#include "Session.h"

#include <async_simple/executors/SimpleExecutor.h>

#include <boost/asio/read.hpp>
#include <boost/asio/write.hpp>
#include <string>

void Session::start() {
    reset_heartbeat_timer();
    boost::asio::co_spawn(socket_.get_executor(), handle(), boost::asio::detached);
}

boost::asio::awaitable<void> Session::handle() {
    try {
        while (true) {
            switch (state_) {
                case State::ReadHeader: {
                    co_await read_header();
                    break;
                }
                case State::ReadLength: {
                    co_await read_length();
                    break;
                }
                case State::ReadBody: {
                    co_await read_body();
                    break;
                }
                case State::Closed: {
                    socket_.close();
                    heartbeat_timer_.cancel();
                    co_return;
                }
            }
        }
    } catch (const boost::system::system_error &e) {
        socket_.close();
        heartbeat_timer_.cancel();
        state_ = State::Closed;
        co_return;
    }
}

boost::asio::awaitable<void> Session::read_header() {
    try {
        co_await boost::asio::async_read(
                socket_, boost::asio::buffer(&header, sizeof(header)),
                boost::asio::use_awaitable);
        reset_heartbeat_timer();

        if (!header) {
            state_ = State::ReadLength;
        } else {
            // 如果是心跳包，则重置为读取头
            state_ = State::ReadHeader;
        }
    } catch (const boost::system::system_error &e) {
        state_ = State::Closed;
    }
    co_return;
}

boost::asio::awaitable<void> Session::read_length() {
    try {
        co_await boost::asio::async_read(
                socket_, boost::asio::buffer(&message_length, sizeof(message_length)),
                boost::asio::use_awaitable);
        message_length = ntohl(message_length);
        state_ = State::ReadBody;
    } catch (const boost::system::system_error &e) {
        state_ = State::Closed;
    }
    co_return;
}

boost::asio::awaitable<void> Session::read_body() {
    try {
        buffer_.resize(message_length);
        co_await boost::asio::async_read(
                socket_, boost::asio::buffer(buffer_, message_length),
                boost::asio::use_awaitable);

        std::string message(buffer_.data(), message_length);

        auto result = kv_middleware_->recv_msg(message);

        co_await write(result);
        state_ = State::ReadHeader;
    } catch (const boost::system::system_error &e) {
        state_ = State::Closed;
    }
    co_return;
}

boost::asio::awaitable<void> Session::write(const std::string &message) {
    try {
        std::vector<char> bufferToSend;
        uint32_t length = htonl(message.length());
        bufferToSend.insert(bufferToSend.end(), reinterpret_cast<char *>(&length),
                            reinterpret_cast<char *>(&length) + sizeof(length));
        bufferToSend.insert(bufferToSend.end(), message.begin(), message.end());
        co_await boost::asio::async_write(socket_, boost::asio::buffer(bufferToSend),
                                          boost::asio::use_awaitable);
        state_ = State::ReadHeader;
    } catch (const boost::system::system_error &e) {
        state_ = State::Closed;
    }
    co_return;
}

void Session::reset_heartbeat_timer() {
    auto self = shared_from_this();
    heartbeat_timer_.expires_after(heartbeat_interval_);
    heartbeat_timer_.async_wait([self](boost::system::error_code ec) {
        if (!ec) {
            // 关闭 socket和计时器
            self->socket_.close();
            self->heartbeat_timer_.cancel();
        }
    });
}